package org.example;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.InetAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class Main {

    static ExecutorService executorService = Executors.newCachedThreadPool();

    public static void main(String[] args) throws IOException {


        try (var serverSocket = new ServerSocket(8888)) {
            System.out.println("Socks5 代理服务器已启动，监听端口: 8888");
            while (!Thread.interrupted()) {
                var socket = serverSocket.accept();
                executorService.submit(() -> {
                    clientHandler(socket);
                });
            }
        }
    }

    private static void clientHandler(Socket socket) {

        try (socket) {

            negotiatedConnection(socket);

            connectServer(socket);

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private static void negotiatedConnection(Socket socket) throws IOException {
        //按需读取数据
        InputStream inputStream = socket.getInputStream();
        //协议判断
        if (inputStream.read() != 5) {
            throw new RuntimeException("只支持socks5协议!!");
        }
        //认证方式
        int authMethodCount = inputStream.read();
        boolean noAuth = false;
        for (int i = 0; i < authMethodCount; i++) {
            int data = inputStream.read();
            if (data == 0) {
                noAuth = true;
            }
        }
        if (!noAuth) {
            throw new RuntimeException("只支持无认证方式!!");
        }
        //响应无需验证
        socket.getOutputStream().write(new byte[]{5, 0});
    }

    private static void connectServer(Socket socket) throws IOException {

        InputStream inputStream = socket.getInputStream();

        int ver = inputStream.read();
        int cmd = inputStream.read();
        int rsv = inputStream.read();
        int atyp = inputStream.read();
        String host = "";

        //用于回复:成功
        var replyStream  = new ByteArrayOutputStream();
        replyStream.write(ver);
        replyStream.write(0);
        replyStream.write(rsv);
        replyStream.write(atyp);

        if (atyp == 1) {//IPv4 地址
            byte[] data = inputStream.readNBytes(4);
            host = InetAddress.getByAddress(data).getHostAddress();
            replyStream.write(data);
        } else if (atyp == 3) {//表示域名
            int len = inputStream.read();
            byte[] data = inputStream.readNBytes(len);
            host = new String(data);
            replyStream.write(len);
            replyStream.write(data);
        } else if (atyp == 4) {//IPv6 地址
            byte[] data = inputStream.readNBytes(16);
            host = InetAddress.getByAddress(data).getHostAddress();
            replyStream.write(data);
        } else {
            throw new RuntimeException("不支持的地址类型: " + atyp);
        }
        byte[] data = inputStream.readNBytes(2);
        int port = ((data[0] & 0xFF) << 8) | (data[1] & 0xFF);
        replyStream.write(data);
        System.out.println("请求连接: " + host + ":" + port);

        Socket targetSocket = new Socket(host, port);
        //响应连接成功
        socket.getOutputStream().write(replyStream.toByteArray());

        executorService.submit(() -> {
            //转发数据
            try (socket) {
                targetSocket.getInputStream().transferTo(socket.getOutputStream());
            } catch (Exception e) {
//                e.printStackTrace();
            }

        });
        //转发数据
        try (targetSocket){
            socket.getInputStream().transferTo(targetSocket.getOutputStream());
        }catch (Exception e) {
//            e.printStackTrace();

        }

    }

}